#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
创建包含 AIGC 元数据的图片文件
支持两种方式：
1. 使用 exiftool 格式的 comment
2. 使用 pyexiv2 嵌入 XMP 元数据
"""

import json
import os
import sys
from datetime import datetime

# 设置输出编码为 UTF-8
if sys.platform == 'win32':
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

try:
    from PIL import Image, ImageDraw, ImageFont
except ImportError:
    print("错误: 需要安装 Pillow")
    print("安装命令: pip install Pillow")
    sys.exit(1)

try:
    from pyexiv2 import Image as ExivImage
    HAS_PYEXIV2 = True
except ImportError:
    print("警告: pyexiv2 未安装，将只使用 comment 方式")
    print("安装命令: pip install pyexiv2")
    HAS_PYEXIV2 = False


def create_test_image(output_path, width=800, height=600):
    """
    创建一个测试图片（PNG 格式）
    """
    # 创建一个彩色渐变图片
    img = Image.new('RGB', (width, height), color='white')
    draw = ImageDraw.Draw(img)
    
    # 绘制渐变背景
    for y in range(height):
        r = int(255 * (y / height))
        g = int(255 * ((height - y) / height))
        b = 128
        draw.line([(0, y), (width, y)], fill=(r, g, b))
    
    # 添加文字
    try:
        # 尝试使用系统字体
        font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", 40)
    except:
        try:
            font = ImageFont.truetype("arial.ttf", 40)
        except:
            font = ImageFont.load_default()
    
    text = "AIGC Test Image"
    bbox = draw.textbbox((0, 0), text, font=font)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]
    
    # 居中显示文字
    x = (width - text_width) // 2
    y = (height - text_height) // 2
    draw.text((x, y), text, fill=(255, 255, 255), font=font)
    
    # 保存为 PNG 格式
    img.save(output_path, 'PNG')
    print(f"[OK] PNG 图片已创建: {output_path}")


def embed_png_metadata(image_path, aigc_data):
    """
    使用 PNG tEXt chunk 嵌入元数据（PNG 格式）
    """
    try:
        from PIL import Image, PngImagePlugin
        
        # 读取图片
        img = Image.open(image_path)
        
        # 创建 comment 字符串
        comment = f'aigc:{json.dumps(aigc_data)}'
        
        # PNG 使用 tEXt chunks 存储元数据
        # 创建 PNG 元数据
        pnginfo = PngImagePlugin.PngInfo()
        pnginfo.add_text("Comment", comment)
        pnginfo.add_text("AIGC", json.dumps(aigc_data))
        
        # 保存图片（PNG 格式）
        img.save(image_path, 'PNG', pnginfo=pnginfo)
        print(f"[OK] PNG tEXt 元数据已嵌入")
        
    except Exception as e:
        print(f"[WARN] PNG 元数据嵌入失败: {e}")


def embed_xmp_metadata(image_path, aigc_data):
    """
    使用 pyexiv2 嵌入 XMP 元数据
    """
    if not HAS_PYEXIV2:
        print("[WARN] pyexiv2 未安装，跳过 XMP 嵌入")
        return
    
    try:
        img = ExivImage(image_path)
        
        aigc_json = json.dumps(aigc_data)
        
        # 方式1: 使用 Xmp.dc.aigc (Dublin Core 命名空间，已存在)
        try:
            img.modify_xmp({
                'Xmp.dc.aigc': aigc_json
            })
            print(f"[OK] XMP.dc.aigc 已设置")
        except Exception as e:
            print(f"[WARN] XMP.dc.aigc 设置失败: {e}")
        
        # 方式2: 使用 Xmp.dc.description 存储 JSON（备用方案）
        try:
            img.modify_xmp({
                'Xmp.dc.description': 'AIGC: ' + aigc_json
            })
        except:
            pass
        
        # 方式3: 使用 Xmp.xmp.Rating 或其他已存在的字段存储 JSON（作为备用）
        # 注意：pyexiv2 需要预先注册自定义命名空间，这里使用标准命名空间
        
        img.close()
        print(f"[OK] XMP 元数据已嵌入")
        
    except Exception as e:
        print(f"[WARN] XMP 嵌入失败: {e}")


def main():
    """
    主函数
    """
    # AIGC 元数据
    aigc_data = {
        "GeneratingTool": "tools",
        "Timestamp": "2021-10-19T00:00:00",
        "ContentID": "V028749274"
    }
    
    # 输出路径
    output_dir = os.path.expanduser("~/Documents/test")
    os.makedirs(output_dir, exist_ok=True)
    
    image_path = os.path.join(output_dir, "testPng.png")
    
    print("=" * 60)
    print("创建包含 AIGC 元数据的图片")
    print("=" * 60)
    print(f"\nAIGC 元数据:")
    print(json.dumps(aigc_data, indent=2, ensure_ascii=False))
    print(f"\n输出路径: {image_path}")
    print()
    
    # 步骤1: 创建测试图片
    print("步骤1: 创建 PNG 测试图片...")
    create_test_image(image_path)
    
    # 步骤2: 嵌入 PNG tEXt 元数据
    print("\n步骤2: 嵌入 PNG tEXt 元数据...")
    embed_png_metadata(image_path, aigc_data)
    
    # 步骤3: 嵌入 XMP 元数据（使用 pyexiv2）
    print("\n步骤3: 嵌入 XMP 元数据...")
    embed_xmp_metadata(image_path, aigc_data)
    
    print("\n" + "=" * 60)
    print("完成！")
    print("=" * 60)
    print(f"\n图片已创建: {image_path}")
    print("\n验证命令:")
    print(f"  exiftool {image_path}")
    print(f"  exiftool -Comment {image_path}")
    print(f"  exiftool -XMP:dc:aigc {image_path}")
    print(f"  python verify_aigc_metadata.py")


if __name__ == "__main__":
    main()

